import logging
import random
import string
from typing import List, Dict
from conf.config import Config
from elasticsearch import AsyncElasticsearch
from elasticsearch.helpers import async_bulk

es_log = logging.getLogger('elasticsearch')
es_log.setLevel(logging.ERROR)


class EsClient:
    def __init__(self, index):
        self.index = index
        self.__host = Config.es.host
        self.__user = Config.es.user
        self.__password = Config.es.password
        self.__port = Config.es.port
        self.es = AsyncElasticsearch(
            hosts=[f"http://{self.__host}:{self.__port}"], basic_auth=(self.__user, self.__password)
        )

    async def __search(self, *args, **kwargs):
        res = await self.es.search(*args, **kwargs)
        await self.es.close()
        return res

    # 设置索引
    def set_index(self, index: str):
        self.index = index

    # 创建索引
    async def create_index(self, index_name: str, mappings=None):
        res = await self.es.indices.create(index=index_name, mappings=mappings)
        await self.es.close()
        return res

    # 删除索引
    async def delete_index(self, index_name: str):
        res = await self.es.indices.delete(index=index_name)
        await self.es.close()
        return res

    # 获取索引
    async def get_index(self, index_name: str):
        res = await self.es.indices.get(index=index_name)
        await self.es.close()
        return res

    # 创建文档（单个）
    async def create_doc(
            self, body,
            _id=''.join(random.sample(string.ascii_letters + string.ascii_uppercase + string.digits, 20))
    ):
        res = await self.es.create(index=self.index, body=body, id=_id)
        await self.es.close()
        return res

    # 创建文档(批量)
    async def create_doc_bulk(self, docs: List[Dict]):
        actions = []
        for doc in docs:
            action = {
                "_index": self.index,
                "_op_type": "create",
                "_id": ''.join(random.sample(string.ascii_letters + string.ascii_uppercase + string.digits, 20))
            }
            for k, v in doc.items():
                action[k] = v
            actions.append(action)
        res = await async_bulk(client=self.es, actions=actions)
        await self.es.close()
        return res

    # 删除文档
    async def delete_doc(self, doc_id):
        res = await self.es.delete(index=self.index, id=doc_id)
        await self.es.close()
        return res

    # 更新文档
    async def update_doc(self, doc_id, doc: Dict):
        body = {
            "doc": doc
        }
        res = await self.es.update(index=self.index, id=doc_id, body=body)
        await self.es.close()
        return res

    # 通过scroll_id分页获取后序文档
    async def get_doc_by_scroll_id(self, scroll_id):
        page = await self.es.scroll(scroll_id=scroll_id, scroll="5m")
        await self.es.close()
        data_list = []
        scroll_id = page.get('_scroll_id')
        for data in page.get('hits').get('hits'):
            data_list.append(data)
        return scroll_id, data_list

    # 清空scroll_id，防止服务端不够用
    async def clear_scroll(self, scroll_id):
        await self.es.clear_scroll(scroll_id)
        await self.es.close()

    # 获取索引的hits内容（一般用于获取文档id、总数）
    async def get_doc_all(self):
        res = await self.__search(index=self.index)
        return res['hits']

    # 获取一个文档
    async def get_doc_by_id(self, id_):
        res = await self.es.get(index=self.index, id=id_)
        await self.es.close()
        return res["_source"]

    # 分页获取超过10000的文档
    async def get_doc_scroll(self, query: dict):
        res = await self.__search(index=self.index, size=10000, body=query, search_type="query_then_fetch", scroll="5m")
        data_list = []
        hits = res.get("hits")
        scroll_id = res.get('_scroll_id')
        total_value = 0
        # total 可能为Dict或int
        if isinstance(hits.get('total'), dict):
            total_value = hits.get('total').get('value')
        else:
            total_value = hits.get('total')

        if total_value > 0:
            for data in hits.get('hits'):
                data_list.append(data.get('_source'))
        return scroll_id, data_list

    # 获取所有文档的_source内容(小于10000)
    async def get_doc(self, query: dict, size: int = 10000):
        query['size'] = size
        res = await self.__search(index=self.index, body=query)
        data_list = []
        hits = res.get("hits")
        # total 可能为Dict或int
        if isinstance(hits.get('total'), dict):
            total_value = hits.get('total').get('value')
        else:
            total_value = hits.get('total')

        if total_value > 0:
            for data in hits.get('hits'):
                data_list.append(data.get('_source'))
        return data_list

    # 聚合查询（分组条件名为group_by, 返回buckets）
    async def get_doc_agg(self, query):
        res = await self.__search(index=self.index, body=query)
        return res['aggregations']['group_by'].get('buckets')

    # 统计查询（统计条件为stats_by，返回最值、平均值等）
    async def get_doc_stats(self, query):
        res = await self.__search(index=self.index, body=query)
        return res['aggregations']["stats_by"]